#include "structure_gst.h"

StructureGST::StructureGST( DataSet *theData ) : AnalysisBase(theData) {
}

bool StructureGST::run() {
    _theData->addToLog("Staring Gst");
    if( !_theData->numInds() ) return false;

    _results->appendHeader("G<sub>ST</sub> Results");

    QStringList locusNames = _theData->getIndividual(0)->locusNames();
    AlleleFrequencies* totalFreqs = _theData->getTotalFrequencies();

    foreach(QString locus, locusNames ) {
        QStringList strata = _theData->getStrataNames();
        qreal Ht = totalFreqs->getLocusFrequency(locus)->getHe();
        qreal Hs = 0.0;
        foreach( QString stratum, strata ){
            AlleleFrequencies* popFreqs = _theData->getStratumFrequencies( stratum );
            Hs += popFreqs->getLocusFrequency( locus )->getHe();
        }
        _Ht[locus] = Ht;
        _Hs[locus] = Hs;
        _Gst[locus] = (Ht > 0 ) ? 1-Hs/Ht : 0.0;
    }
    _theData->addToLog( "Finished Gst");
    return true;
}


qreal StructureGST::getGSTForLocus( QString key ) {
    return ( _Gst.keys().contains( key )) ? _Gst[key] : 0.0;
}

qreal StructureGST::getHsForLocus( QString key ) {
    return ( _Hs.keys().contains( key )) ? _Gst[key] : 0.0;
}

qreal StructureGST::getHtForLocus( QString key ) {
    return ( _Ht.keys().contains( key )) ? _Gst[key] : 0.0;
}

qreal StructureGST::getMultilocusGst() {
    qreal Ht = 0.0;
    qreal Hs = 0.0;
    foreach(QString key, _Hs.keys()) {
        Ht += _Ht[key];
        Hs += _Hs[key];
    }
    return (Ht>0) ? 1-Hs/Ht : 0.0;
}
